# 1.使用sklearn配合pytorch完成以下操作（每题10分）
# (1)数据处理
import matplotlib.pyplot as plt
import torch
from sklearn.datasets import load_iris

torch.manual_seed(1234)
# ①读取iris数据集
x, y = load_iris(return_X_y=True)
# ②自定义超参数数值
# ③将数据7:3切分
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(x, y, train_size=0.7)
# ④放入dataloder
from torch.utils.data import DataLoader, TensorDataset

train_dataset = TensorDataset(torch.tensor(X_train), torch.tensor(y_train))
test_dataset = TensorDataset(torch.tensor(X_test), torch.tensor(y_test))
train_loader = DataLoader(train_dataset, batch_size=10)
test_loader = DataLoader(test_dataset, batch_size=10)
# ⑤创建全连接层，两个隐藏层，神经元数量分别是10,5
# ⑥每个隐藏层后用relu激活，dropout设定0.2

model = torch.nn.Sequential(
    torch.nn.Linear(in_features=4, out_features=10),
    torch.nn.ReLU(),
    torch.nn.Dropout(0.2),
    torch.nn.Linear(in_features=10, out_features=5),
    torch.nn.ReLU(),
    torch.nn.Dropout(0.2),
    torch.nn.Linear(in_features=5, out_features=3),
)
# ⑦使用交叉熵作为代价函数
loss_fn = torch.nn.CrossEntropyLoss()
op = torch.optim.Adagrad(model.parameters(), lr=0.1)
print()
train_loss_list = []
test_loss_list = []
# ⑧训练模型，每100批次打印损失值和准确率

for i in range(300):
    model.train()
    train_loss = 0
    for j, (x, y) in enumerate(train_loader):
        # torch.Tensor.float()
        op.zero_grad()
        predict = model(x.float())
        loss = loss_fn(predict, y)
        loss.backward()
        op.step()
        train_loss += loss.item()
    train_loss = train_loss / len(train_loader)
    train_loss_list.append(train_loss)

    model.eval()
    test_loss = 0
    for (x, y) in test_loader:
        predict = model(x.float())
        loss = loss_fn(predict, y)
        test_loss += loss.item()
    test_loss = test_loss / len(test_loader)
    test_loss_list.append(test_loss)

# ⑨绘制损失变化曲线
plt.plot(train_loss_list, c='r')
plt.plot(test_loss_list, c='g')
plt.show()
# ⑩打印最终准确率
avg_acc = 0
for (x, y) in test_loader:
    predict = model(x.float())
    avg_acc += predict.argmax(1).eq(y).float().sum().item()
print('acc is', avg_acc / len(test_dataset))
